import os

# 设置HF_ENDPOINT环境变量
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
from sentence_transformers import SentenceTransformer, util
import torch
import pprint

def start_model():
    print('start text2vec model')

    # 检查是否有可用的GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 加载模型，并将其移动到相应的设备上
    m = SentenceTransformer("moka-ai/m3e-base").to(device)

    # 设置模型为评估模式
    m.eval()

    print('finish')
    return m


def rank_docs(m, query, docs, docs_embedding, topk,return_idx = True):
    query_embedding = m.encode([query], convert_to_tensor=True)
    cosine_scores = util.pytorch_cos_sim(query_embedding, docs_embedding)[0]
    if return_idx:
        return cosine_scores.argsort(descending=True).tolist()
    sorted_sentences = [docs[i] for i in cosine_scores.argsort(descending=True)]
    # 输出排序后的句子
    print('topk',topk)
    return sorted_sentences[:topk]


class Documents:
    def __init__(self,model,  sources,documents, topk=3):
        print('init documents class')
        self.model = model
        self.sources = sources
        self.id2sources = {i: s for i, s in enumerate(sources)}

        self.documents = documents

        self.topk = topk
        self.embeddings = self.build_embedding(documents)

    def build_embedding(self, documents):
        res = self.model.encode(documents, convert_to_tensor=True)
        return res

    def rank_docs(self, query,topk=10):
        print('排序结果')
        res = rank_docs(self.model, query, self.documents, self.embeddings,topk)
        return res

    def rank_docs_with_condition(self, query,filter:dict , topk=10):
        print('排序结果')
        res = rank_docs(self.model, query, self.documents, self.embeddings,topk)
        def filter_f(d):
            return all(d[k] == v for k, v in filter.items())
        rank_res = [self.id2sources[i] for i in res]
        pprint.pprint({'rank_res':rank_res})
        source_res = [e for e in rank_res if filter_f(e)  ]
        pprint.pprint({'source_res': source_res})
        return source_res[:topk]

    def get_chat_prompt(self, query):
        docs = self.get_similarities(query)
        new_docs = []
        for d in docs:
            if d not in new_docs:
                new_docs.append(d)
        docs = new_docs[:self.topk]

        docs_content_t = '\n\n'.join(new_docs)
        info1 = f'根据以下信息回答问题,用简洁清晰的话回答，并且不要提及无关的内容：\n\n已知信息：\n{docs_content_t}\n\n问题：'
        info2 = f'根据以下信息回答问题,用简洁清晰的话回答，并且不要提及无关的内容：\n\n已知信息：\n{docs_content_t}\n\n问题：{query}'
        return info1, info2, docs


#
# doc_list = ["我喜欢听歌","我喜欢打篮球","我经常用耳机"]
# DocDb = Documents(start_model(),doc_list)
# print(DocDb.get_similarities("我喜欢乔丹"))

data = {
    "query": "what is ThinkAny",
    "search_results": [
        {
            "title": "ThinkAny AI - LinkedIn",
            "link": "https://hk.linkedin.com/in/thinkany-ai-3bb267300",
            "snippet": "Experience: ThinkAny · Location: Hong Kong. View ThinkAny AI's profile on LinkedIn, a professional community of 1 billion members.",
            "position": 1,
            "uuid": "88c7c8b192437cb3025bf3afd3ce6fa7"
        },
        {
            "title": "ThinkAny - AI Search Engine",
            "link": "https://thinkany.ai/",
            "snippet": "ThinkAny is a free AI search engine using advanced RAG vector search for precise, trusted answers, with interactive AI assistant chat for comprehensive user ...",
            "position": 2,
            "uuid": "5209e9df294fc9d8fdf20a023b1a561c"
        },
        {
            "title": "thinkany-ai - GitHub",
            "link": "https://github.com/thinkany-ai",
            "snippet": "Skip to content. Navigation Menu. Toggle navigation. Sign in · thinkany-ai. Product. Actions. Automate any workflow · Packages. Host and manage packages.",
            "position": 3,
            "uuid": "f20e53a4367b053f4eb88fb462670e1b"
        },
        {
            "title": "ThinkAny - AI Search Engine And 39 Other AI Alternatives For ...",
            "link": "https://theresanaiforthat.com/ai/thinkany-ai-search-engine/",
            "snippet": "ThinkAny is an advanced AI-powered search engine bringing together quality content retrieval and innovative AI-driven answering capabilities ...",
            "position": 4,
            "uuid": "c23fbd932565c797c907687794fe412b"
        },
        {
            "title": "ThinkAny: The AI Search Engine Revolution - Supertools",
            "link": "https://supertools.therundown.ai/content/thinkany",
            "snippet": "ThinkAny is an advanced AI search engine that combines RAG technology with intelligent answering capabilities to deliver high-quality content and efficient ...",
            "position": 5,
            "uuid": "d5c608a64e72e76b0910aabf18ef3a32"
        }
    ]

}


def print_list(es):
    for e in es:
        print(e)


samples = [
    {'remark_name':'123','总结':'我喜欢吃水果'},
    {'remark_name':'123','总结':'我喜欢吃苹果'},
    {'remark_name':'456','总结':'我喜欢游泳'},

           ]


doc_list = ["我喜欢听歌", "我喜欢打篮球", "我经常用耳机",]

docs = [(d['position'], d['title'], d['snippet']) for d in data['search_results']]
m = start_model()

def run_on_samples():
    import pandas as pd
    df = pd.read_csv('samples.csv')
    df['remark_name'] = df['remark_name'].astype(str)
    samples = df.to_dict('records')
    doc_list = [s['总结'] for s in  samples ]#
    # 构建向量， 全量构建
    DocDb = Documents(m, samples,doc_list)
    # 做一次查询
    res2 = DocDb.rank_docs_with_condition("我喜欢苹果",filter={'remark_name':'123'},topk=3)
    print_list(res2)



run_on_samples()
def run():
    doc_list = [title for ix, title, snip in docs]
    doc_list2 = [snip for ix, title, snip in docs]
    doc_list3 = [title + snip for ix, title, snip in docs]

    print('\n\ntitle')
    DocDb = Documents(m, doc_list)
    res = DocDb.rank_docs("what is ThinkAny")
    print_list(res)

    print('\n\nsnip')
    DocDb = Documents(m, doc_list2)

    res = DocDb.rank_docs("what is ThinkAny")
    print_list(res)
    print('\n\n title+snip')
    DocDb = Documents(m, doc_list3)

    res = DocDb.rank_docs("what is ThinkAny")
    print_list(res)